package com.ptp.backendcertms.networksecurity.filter;

import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;
import java.nio.charset.StandardCharsets;

/**
 * @Name ResponseWrapper 类
 * @Author songcongcong
 * @Date 2022/3/9
 * @Description 继承HttpServletResponseWrapper，重写响应，增加或减少响应参数
 * 获取HttpServletResponse的响应需要先获取HttpServletResponse的输出流，为避免获取一次之后，再次获取报Stream Closed异常，暂存输出流
 */
public class ResponseWrapper extends HttpServletResponseWrapper {
    private ByteArrayOutputStream buffer;
    private ServletOutputStream out;
    private PrintWriter writer;

    public ResponseWrapper(HttpServletResponse response) throws IOException {
        super(response);
        //真正存储数据的流
        buffer = new ByteArrayOutputStream();
        out = new WrapperOutputStream(buffer);
        writer = new PrintWriter(new OutputStreamWriter(buffer, StandardCharsets.UTF_8));
    }

    /**
     * 重载父类获取outputstream的方法
     *
     * @return
     * @throws IOException
     */
    @Override
    public ServletOutputStream getOutputStream() {
        return out;
    }

    /**
     * 重载父类获取writer的方法
     *
     * @return
     * @throws UnsupportedEncodingException
     */
    @Override
    public PrintWriter getWriter() {
        return writer;
    }

    /**
     * 重载父类获取flushBuffer的方法
     *
     * @throws IOException
     */
    @Override
    public void flushBuffer() throws IOException {
        if (out != null) {
            out.flush();
        }
        if (writer != null) {
            writer.flush();
        }
    }

    @Override
    public void reset() {
        buffer.reset();
    }

    /**
     * 将out、writer中的数据强制输出到WapperedResponse的buffer里面，否则取不到数据
     *
     * @return
     * @throws IOException
     */
    public byte[] getResponseData() throws IOException {
        flushBuffer();
        return buffer.toByteArray();
    }
}
